import sys
from io import BytesIO, StringIO
from threading import Thread

from paramiko.channel import Channel, ChannelFile, ChannelStderrFile, ChannelStdinFile

from .utils import threaded_io_transfer


class SshCommand:
    def __init__(self, channel: Channel, command: str):
        self._channel = channel
        self._command = command
        self._stdout_thread: Thread | None = None
        self._stderr_thread: Thread | None = None
        self._exit_code: int | None = None

    @property
    def command(self) -> str:
        return self._command

    @property
    def exit_code(self) -> int | None:
        return self._exit_code

    def stdin(self, *, buffer_size: int = -1) -> ChannelStdinFile:
        return self._channel.makefile_stdin("wb", buffer_size)

    def stdout(self, *, buffer_size: int = -1) -> ChannelFile:
        return self._channel.makefile("r", buffer_size)

    def stderr(self, *, buffer_size: int = -1) -> ChannelStderrFile:
        return self._channel.makefile_stderr("r", buffer_size)

    def attach_stdout(self, io: BytesIO | StringIO | None = None, *, buffer_size: int = -1):
        if self._stdout_thread is not None:
            raise RuntimeError("Cannot attach stdout twice!")

        # noinspection PyTypeChecker
        in_io: BytesIO = self.stdout(buffer_size=buffer_size)
        out_io = io if io is not None else sys.stdout
        self._stdout_thread = threaded_io_transfer(in_io, out_io, by_line=True)

    def attach_stderr(self, io: BytesIO | StringIO | None = None, *, buffer_size: int = -1):
        if self._stderr_thread is not None:
            raise RuntimeError("Cannot attach stderr twice!")

        # noinspection PyTypeChecker
        in_io: BytesIO = self.stderr(buffer_size=buffer_size)
        out_io = io if io is not None else sys.stderr
        self._stderr_thread = threaded_io_transfer(in_io, out_io, by_line=True)

    def execute(self):
        self._channel.exec_command(self._command)

    def interrupt(self):
        self._channel.close()
        if self._stdout_thread is not None:
            self._stdout_thread.join()
        if self._stderr_thread is not None:
            self._stderr_thread.join()

    def wait(self, *, raise_error: bool = True):
        code = self._channel.recv_exit_status()
        if self._stdout_thread is not None:
            self._stdout_thread.join()
        if self._stderr_thread is not None:
            self._stderr_thread.join()
        if raise_error and code != 0:
            raise RuntimeError(f"Command exited with non-zero code: {code}")
        self._exit_code = code
